{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Surname Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from argparse import Namespace\n",
    "from collections import Counter\n",
    "import json\n",
    "import re\n",
    "import string\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm_notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Vectorization classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocabulary(object):\n",
    "    \"\"\"Class to process text and extract vocabulary for mapping\"\"\"\n",
    "\n",
    "    def __init__(self, token_to_idx=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            token_to_idx (dict): a pre-existing map of tokens to indices\n",
    "        \"\"\"\n",
    "\n",
    "        if token_to_idx is None:\n",
    "            token_to_idx = {}\n",
    "        self._token_to_idx = token_to_idx\n",
    "\n",
    "        self._idx_to_token = {idx: token \n",
    "                              for token, idx in self._token_to_idx.items()}\n",
    "        \n",
    "    def to_serializable(self):\n",
    "        \"\"\" returns a dictionary that can be serialized \"\"\"\n",
    "        return {'token_to_idx': self._token_to_idx}\n",
    "\n",
    "    @classmethod\n",
    "    def from_serializable(cls, contents):\n",
    "        \"\"\" instantiates the Vocabulary from a serialized dictionary \"\"\"\n",
    "        return cls(**contents)\n",
    "\n",
    "    def add_token(self, token):\n",
    "        \"\"\"Update mapping dicts based on the token.\n",
    "\n",
    "        Args:\n",
    "            token (str): the item to add into the Vocabulary\n",
    "        Returns:\n",
    "            index (int): the integer corresponding to the token\n",
    "        \"\"\"\n",
    "        if token in self._token_to_idx:\n",
    "            index = self._token_to_idx[token]\n",
    "        else:\n",
    "            index = len(self._token_to_idx)\n",
    "            self._token_to_idx[token] = index\n",
    "            self._idx_to_token[index] = token\n",
    "        return index\n",
    "            \n",
    "    def add_many(self, tokens):\n",
    "        \"\"\"Add a list of tokens into the Vocabulary\n",
    "        \n",
    "        Args:\n",
    "            tokens (list): a list of string tokens\n",
    "        Returns:\n",
    "            indices (list): a list of indices corresponding to the tokens\n",
    "        \"\"\"\n",
    "        return [self.add_token(token) for token in tokens]\n",
    "\n",
    "    def lookup_token(self, token):\n",
    "        \"\"\"Retrieve the index associated with the token \n",
    "        \n",
    "        Args:\n",
    "            token (str): the token to look up \n",
    "        Returns:\n",
    "            index (int): the index corresponding to the token\n",
    "        \"\"\"\n",
    "        return self._token_to_idx[token]\n",
    "\n",
    "    def lookup_index(self, index):\n",
    "        \"\"\"Return the token associated with the index\n",
    "        \n",
    "        Args: \n",
    "            index (int): the index to look up\n",
    "        Returns:\n",
    "            token (str): the token corresponding to the index\n",
    "        Raises:\n",
    "            KeyError: if the index is not in the Vocabulary\n",
    "        \"\"\"\n",
    "        if index not in self._idx_to_token:\n",
    "            raise KeyError(\"the index (%d) is not in the Vocabulary\" % index)\n",
    "        return self._idx_to_token[index]\n",
    "\n",
    "    def __str__(self):\n",
    "        return \"<Vocabulary(size=%d)>\" % len(self)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._token_to_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequenceVocabulary(Vocabulary):\n",
    "    def __init__(self, token_to_idx=None, unk_token=\"<UNK>\",\n",
    "                 mask_token=\"<MASK>\", begin_seq_token=\"<BEGIN>\",\n",
    "                 end_seq_token=\"<END>\"):\n",
    "\n",
    "        super(SequenceVocabulary, self).__init__(token_to_idx)\n",
    "\n",
    "        self._mask_token = mask_token\n",
    "        self._unk_token = unk_token\n",
    "        self._begin_seq_token = begin_seq_token\n",
    "        self._end_seq_token = end_seq_token\n",
    "\n",
    "        self.mask_index = self.add_token(self._mask_token)\n",
    "        self.unk_index = self.add_token(self._unk_token)\n",
    "        self.begin_seq_index = self.add_token(self._begin_seq_token)\n",
    "        self.end_seq_index = self.add_token(self._end_seq_token)\n",
    "\n",
    "    def to_serializable(self):\n",
    "        contents = super(SequenceVocabulary, self).to_serializable()\n",
    "        contents.update({'unk_token': self._unk_token,\n",
    "                         'mask_token': self._mask_token,\n",
    "                         'begin_seq_token': self._begin_seq_token,\n",
    "                         'end_seq_token': self._end_seq_token})\n",
    "        return contents\n",
    "\n",
    "    def lookup_token(self, token):\n",
    "        \"\"\"Retrieve the index associated with the token \n",
    "          or the UNK index if token isn't present.\n",
    "        \n",
    "        Args:\n",
    "            token (str): the token to look up \n",
    "        Returns:\n",
    "            index (int): the index corresponding to the token\n",
    "        Notes:\n",
    "            `unk_index` needs to be >=0 (having been added into the Vocabulary) \n",
    "              for the UNK functionality \n",
    "        \"\"\"\n",
    "        if self.unk_index >= 0:\n",
    "            return self._token_to_idx.get(token, self.unk_index)\n",
    "        else:\n",
    "            return self._token_to_idx[token]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SurnameVectorizer(object):\n",
    "    \"\"\" The Vectorizer which coordinates the Vocabularies and puts them to use\"\"\"    \n",
    "    def __init__(self, char_vocab, nationality_vocab):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            char_vocab (SequenceVocabulary): maps words to integers\n",
    "            nationality_vocab (Vocabulary): maps nationalities to integers\n",
    "        \"\"\"\n",
    "        self.char_vocab = char_vocab\n",
    "        self.nationality_vocab = nationality_vocab\n",
    "\n",
    "    def vectorize(self, surname, vector_length=-1):\n",
    "        \"\"\"Vectorize a surname into a vector of observations and targets\n",
    "        \n",
    "        The outputs are the vectorized surname split into two vectors:\n",
    "            surname[:-1] and surname[1:]\n",
    "        At each timestep, the first vector is the observation and the second vector is the target. \n",
    "        \n",
    "        Args:\n",
    "            surname (str): the surname to be vectorized\n",
    "            vector_length (int): an argument for forcing the length of index vector\n",
    "        Returns:\n",
    "            a tuple: (from_vector, to_vector)\n",
    "            from_vector (numpy.ndarray): the observation vector \n",
    "            to_vector (numpy.ndarray): the target prediction vector\n",
    "        \"\"\"\n",
    "        indices = [self.char_vocab.begin_seq_index] \n",
    "        indices.extend(self.char_vocab.lookup_token(token) for token in surname)\n",
    "        indices.append(self.char_vocab.end_seq_index)\n",
    "\n",
    "        if vector_length < 0:\n",
    "            vector_length = len(indices) - 1\n",
    "\n",
    "        from_vector = np.zeros(vector_length, dtype=np.int64)         \n",
    "        from_indices = indices[:-1]\n",
    "        from_vector[:len(from_indices)] = from_indices\n",
    "        from_vector[len(from_indices):] = self.char_vocab.mask_index\n",
    "\n",
    "        to_vector = np.zeros(vector_length, dtype=np.int64)\n",
    "        to_indices = indices[1:]\n",
    "        to_vector[:len(to_indices)] = to_indices\n",
    "        to_vector[len(to_indices):] = self.char_vocab.mask_index\n",
    "        \n",
    "        return from_vector, to_vector\n",
    "\n",
    "    @classmethod\n",
    "    def from_dataframe(cls, surname_df):\n",
    "        \"\"\"Instantiate the vectorizer from the dataset dataframe\n",
    "        \n",
    "        Args:\n",
    "            surname_df (pandas.DataFrame): the surname dataset\n",
    "        Returns:\n",
    "            an instance of the SurnameVectorizer\n",
    "        \"\"\"\n",
    "        char_vocab = SequenceVocabulary()\n",
    "        nationality_vocab = Vocabulary()\n",
    "\n",
    "        for index, row in surname_df.iterrows():\n",
    "            for char in row.surname:\n",
    "                char_vocab.add_token(char)\n",
    "            nationality_vocab.add_token(row.nationality)\n",
    "\n",
    "        return cls(char_vocab, nationality_vocab)\n",
    "\n",
    "    @classmethod\n",
    "    def from_serializable(cls, contents):\n",
    "        \"\"\"Instantiate the vectorizer from saved contents\n",
    "        \n",
    "        Args:\n",
    "            contents (dict): a dict holding two vocabularies for this vectorizer\n",
    "                This dictionary is created using `vectorizer.to_serializable()`\n",
    "        Returns:\n",
    "            an instance of SurnameVectorizer\n",
    "        \"\"\"\n",
    "        char_vocab = SequenceVocabulary.from_serializable(contents['char_vocab'])\n",
    "        nat_vocab =  Vocabulary.from_serializable(contents['nationality_vocab'])\n",
    "\n",
    "        return cls(char_vocab=char_vocab, nationality_vocab=nat_vocab)\n",
    "\n",
    "    def to_serializable(self):\n",
    "        \"\"\" Returns the serializable contents \"\"\"\n",
    "        return {'char_vocab': self.char_vocab.to_serializable(), \n",
    "                'nationality_vocab': self.nationality_vocab.to_serializable()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SurnameDataset(Dataset):\n",
    "    def __init__(self, surname_df, vectorizer):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            surname_df (pandas.DataFrame): the dataset\n",
    "            vectorizer (SurnameVectorizer): vectorizer instatiated from dataset\n",
    "        \"\"\"\n",
    "        self.surname_df = surname_df \n",
    "        self._vectorizer = vectorizer\n",
    "\n",
    "        self._max_seq_length = max(map(len, self.surname_df.surname)) + 2\n",
    "\n",
    "        self.train_df = self.surname_df[self.surname_df.split=='train']\n",
    "        self.train_size = len(self.train_df)\n",
    "\n",
    "        self.val_df = self.surname_df[self.surname_df.split=='val']\n",
    "        self.validation_size = len(self.val_df)\n",
    "\n",
    "        self.test_df = self.surname_df[self.surname_df.split=='test']\n",
    "        self.test_size = len(self.test_df)\n",
    "\n",
    "        self._lookup_dict = {'train': (self.train_df, self.train_size), \n",
    "                             'val': (self.val_df, self.validation_size), \n",
    "                             'test': (self.test_df, self.test_size)}\n",
    "\n",
    "        self.set_split('train')\n",
    "        \n",
    "    @classmethod\n",
    "    def load_dataset_and_make_vectorizer(cls, surname_csv):\n",
    "        \"\"\"Load dataset and make a new vectorizer from scratch\n",
    "        \n",
    "        Args:\n",
    "            surname_csv (str): location of the dataset\n",
    "        Returns:\n",
    "            an instance of SurnameDataset\n",
    "        \"\"\"\n",
    "        \n",
    "        surname_df = pd.read_csv(surname_csv)\n",
    "        return cls(surname_df, SurnameVectorizer.from_dataframe(surname_df))\n",
    "        \n",
    "    @classmethod\n",
    "    def load_dataset_and_load_vectorizer(cls, surname_csv, vectorizer_filepath):\n",
    "        \"\"\"Load dataset and the corresponding vectorizer. \n",
    "        Used in the case in the vectorizer has been cached for re-use\n",
    "        \n",
    "        Args:\n",
    "            surname_csv (str): location of the dataset\n",
    "            vectorizer_filepath (str): location of the saved vectorizer\n",
    "        Returns:\n",
    "            an instance of SurnameDataset\n",
    "        \"\"\"\n",
    "        surname_df = pd.read_csv(surname_csv)\n",
    "        vectorizer = cls.load_vectorizer_only(vectorizer_filepath)\n",
    "        return cls(surname_df, vectorizer)\n",
    "\n",
    "    @staticmethod\n",
    "    def load_vectorizer_only(vectorizer_filepath):\n",
    "        \"\"\"a static method for loading the vectorizer from file\n",
    "        \n",
    "        Args:\n",
    "            vectorizer_filepath (str): the location of the serialized vectorizer\n",
    "        Returns:\n",
    "            an instance of SurnameVectorizer\n",
    "        \"\"\"\n",
    "        with open(vectorizer_filepath) as fp:\n",
    "            return SurnameVectorizer.from_serializable(json.load(fp))\n",
    "\n",
    "    def save_vectorizer(self, vectorizer_filepath):\n",
    "        \"\"\"saves the vectorizer to disk using json\n",
    "        \n",
    "        Args:\n",
    "            vectorizer_filepath (str): the location to save the vectorizer\n",
    "        \"\"\"\n",
    "        with open(vectorizer_filepath, \"w\") as fp:\n",
    "            json.dump(self._vectorizer.to_serializable(), fp)\n",
    "\n",
    "    def get_vectorizer(self):\n",
    "        \"\"\" returns the vectorizer \"\"\"\n",
    "        return self._vectorizer\n",
    "\n",
    "    def set_split(self, split=\"train\"):\n",
    "        self._target_split = split\n",
    "        self._target_df, self._target_size = self._lookup_dict[split]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self._target_size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"the primary entry point method for PyTorch datasets\n",
    "        \n",
    "        Args:\n",
    "            index (int): the index to the data point \n",
    "        Returns:\n",
    "            a dictionary holding the data point: (x_data, y_target, class_index)\n",
    "        \"\"\"\n",
    "        row = self._target_df.iloc[index]\n",
    "        \n",
    "        from_vector, to_vector = \\\n",
    "            self._vectorizer.vectorize(row.surname, self._max_seq_length)\n",
    "        \n",
    "        nationality_index = \\\n",
    "            self._vectorizer.nationality_vocab.lookup_token(row.nationality)\n",
    "\n",
    "        return {'x_data': from_vector, \n",
    "                'y_target': to_vector, \n",
    "                'class_index': nationality_index}\n",
    "\n",
    "    def get_num_batches(self, batch_size):\n",
    "        \"\"\"Given a batch size, return the number of batches in the dataset\n",
    "        \n",
    "        Args:\n",
    "            batch_size (int)\n",
    "        Returns:\n",
    "            number of batches in the dataset\n",
    "        \"\"\"\n",
    "        return len(self) // batch_size\n",
    "    \n",
    "def generate_batches(dataset, batch_size, shuffle=True,\n",
    "                     drop_last=True, device=\"cpu\"): \n",
    "    \"\"\"\n",
    "    A generator function which wraps the PyTorch DataLoader. It will \n",
    "      ensure each tensor is on the write device location.\n",
    "    \"\"\"\n",
    "    dataloader = DataLoader(dataset=dataset, batch_size=batch_size,\n",
    "                            shuffle=shuffle, drop_last=drop_last)\n",
    "\n",
    "    for data_dict in dataloader:\n",
    "        out_data_dict = {}\n",
    "        for name, tensor in data_dict.items():\n",
    "            out_data_dict[name] = data_dict[name].to(device)\n",
    "        yield out_data_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.ones(5, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = x.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = torch.zeros(5)\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., 0., 0., 0., 0.], device='cuda:0')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.to(x.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Model: SurnameGenerationModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SurnameGenerationModel(nn.Module):\n",
    "    def __init__(self, char_embedding_size, char_vocab_size, num_nationalities,\n",
    "                 rnn_hidden_size, batch_first=True, padding_idx=0, dropout_p=0.5):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            char_embedding_size (int): The size of the character embeddings\n",
    "            char_vocab_size (int): The number of characters to embed\n",
    "            num_nationalities (int): The size of the prediction vector \n",
    "            rnn_hidden_size (int): The size of the RNN's hidden state\n",
    "            batch_first (bool): Informs whether the input tensors will \n",
    "                have batch or the sequence on the 0th dimension\n",
    "            padding_idx (int): The index for the tensor padding; \n",
    "                see torch.nn.Embedding\n",
    "            dropout_p (float): the probability of zeroing activations using\n",
    "                the dropout method.  higher means more likely to zero.\n",
    "        \"\"\"\n",
    "        super(SurnameGenerationModel, self).__init__()\n",
    "        \n",
    "        self.char_emb = nn.Embedding(num_embeddings=char_vocab_size,\n",
    "                                     embedding_dim=char_embedding_size,\n",
    "                                     padding_idx=padding_idx)\n",
    "\n",
    "        self.nation_emb = nn.Embedding(num_embeddings=num_nationalities,\n",
    "                                       embedding_dim=rnn_hidden_size)\n",
    "\n",
    "        self.rnn = nn.GRU(input_size=char_embedding_size, \n",
    "                          hidden_size=rnn_hidden_size,\n",
    "                          batch_first=batch_first)\n",
    "        \n",
    "        self.fc = nn.Linear(in_features=rnn_hidden_size, \n",
    "                            out_features=char_vocab_size)\n",
    "        \n",
    "        self._dropout_p = dropout_p\n",
    "\n",
    "    def forward(self, x_in, nationality_index, apply_softmax=False):\n",
    "        \"\"\"The forward pass of the model\n",
    "        \n",
    "        Args:\n",
    "            x_in (torch.Tensor): an input data tensor. \n",
    "                x_in.shape should be (batch, max_seq_size)\n",
    "            nationality_index (torch.Tensor): The index of the nationality for each data point\n",
    "                Used to initialize the hidden state of the RNN\n",
    "            apply_softmax (bool): a flag for the softmax activation\n",
    "                should be false if used with the Cross Entropy losses\n",
    "        Returns:\n",
    "            the resulting tensor. tensor.shape should be (batch, char_vocab_size)\n",
    "        \"\"\"\n",
    "        x_embedded = self.char_emb(x_in)\n",
    "        \n",
    "        # hidden_size: (num_layers * num_directions, batch_size, rnn_hidden_size)\n",
    "        nationality_embedded = self.nation_emb(nationality_index).unsqueeze(0)\n",
    "\n",
    "        y_out, _ = self.rnn(x_embedded, nationality_embedded)\n",
    "\n",
    "        batch_size, seq_size, feat_size = y_out.shape\n",
    "        y_out = y_out.contiguous().view(batch_size * seq_size, feat_size)\n",
    "\n",
    "        y_out = self.fc(F.dropout(y_out, p=self._dropout_p))\n",
    "                         \n",
    "        if apply_softmax:\n",
    "            y_out = F.softmax(y_out, dim=1)\n",
    "\n",
    "        new_feat_size = y_out.shape[-1]\n",
    "        y_out = y_out.view(batch_size, seq_size, new_feat_size)\n",
    "            \n",
    "        return y_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_from_model(model, vectorizer, nationalities, sample_size=20, \n",
    "                      temperature=1.0):\n",
    "    \"\"\"Sample a sequence of indices from the model\n",
    "    \n",
    "    Args:\n",
    "        model (SurnameGenerationModel): the trained model\n",
    "        vectorizer (SurnameVectorizer): the corresponding vectorizer\n",
    "        nationalities (list): a list of integers representing nationalities\n",
    "        sample_size (int): the max length of the samples\n",
    "        temperature (float): accentuates or flattens \n",
    "            the distribution. \n",
    "            0.0 < temperature < 1.0 will make it peakier. \n",
    "            temperature > 1.0 will make it more uniform\n",
    "    Returns:\n",
    "        indices (torch.Tensor): the matrix of indices; \n",
    "        shape = (num_samples, sample_size)\n",
    "    \"\"\"\n",
    "    num_samples = len(nationalities)\n",
    "    begin_seq_index = [vectorizer.char_vocab.begin_seq_index \n",
    "                       for _ in range(num_samples)]\n",
    "    begin_seq_index = torch.tensor(begin_seq_index, \n",
    "                                   dtype=torch.int64).unsqueeze(dim=1)\n",
    "    indices = [begin_seq_index]\n",
    "    nationality_indices = torch.tensor(nationalities, dtype=torch.int64).unsqueeze(dim=0)\n",
    "    h_t = model.nation_emb(nationality_indices)\n",
    "    \n",
    "    for time_step in range(sample_size):\n",
    "        x_t = indices[time_step]\n",
    "        x_emb_t = model.char_emb(x_t)\n",
    "        rnn_out_t, h_t = model.rnn(x_emb_t, h_t)\n",
    "        prediction_vector = model.fc(rnn_out_t.squeeze(dim=1))\n",
    "        probability_vector = F.softmax(prediction_vector / temperature, dim=1)\n",
    "        indices.append(torch.multinomial(probability_vector, num_samples=1))\n",
    "    indices = torch.stack(indices).squeeze().permute(1, 0)\n",
    "    return indices\n",
    "\n",
    "def decode_samples(sampled_indices, vectorizer):\n",
    "    \"\"\"Transform indices into the string form of a surname\n",
    "    \n",
    "    Args:\n",
    "        sampled_indices (torch.Tensor): the inidces from `sample_from_model`\n",
    "        vectorizer (SurnameVectorizer): the corresponding vectorizer\n",
    "    \"\"\"\n",
    "    decoded_surnames = []\n",
    "    vocab = vectorizer.char_vocab\n",
    "    \n",
    "    for sample_index in range(sampled_indices.shape[0]):\n",
    "        surname = \"\"\n",
    "        for time_step in range(sampled_indices.shape[1]):\n",
    "            sample_item = sampled_indices[sample_index, time_step].item()\n",
    "            if sample_item == vocab.begin_seq_index:\n",
    "                continue\n",
    "            elif sample_item == vocab.end_seq_index:\n",
    "                break\n",
    "            else:\n",
    "                surname += vocab.lookup_index(sample_item)\n",
    "        decoded_surnames.append(surname)\n",
    "    return decoded_surnames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Routine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_train_state(args):\n",
    "    return {'stop_early': False,\n",
    "            'early_stopping_step': 0,\n",
    "            'early_stopping_best_val': 1e8,\n",
    "            'learning_rate': args.learning_rate,\n",
    "            'epoch_index': 0,\n",
    "            'train_loss': [],\n",
    "            'train_acc': [],\n",
    "            'val_loss': [],\n",
    "            'val_acc': [],\n",
    "            'test_loss': -1,\n",
    "            'test_acc': -1,\n",
    "            'model_filename': args.model_state_file}\n",
    "\n",
    "def update_train_state(args, model, train_state):\n",
    "    \"\"\"Handle the training state updates.\n",
    "    Components:\n",
    "     - Early Stopping: Prevent overfitting.\n",
    "     - Model Checkpoint: Model is saved if the model is better\n",
    "    \n",
    "    :param args: main arguments\n",
    "    :param model: model to train\n",
    "    :param train_state: a dictionary representing the training state values\n",
    "    :returns:\n",
    "        a new train_state\n",
    "    \"\"\"\n",
    "\n",
    "    # Save one model at least\n",
    "    if train_state['epoch_index'] == 0:\n",
    "        torch.save(model.state_dict(), train_state['model_filename'])\n",
    "        train_state['stop_early'] = False\n",
    "\n",
    "    # Save model if performance improved\n",
    "    elif train_state['epoch_index'] >= 1:\n",
    "        loss_tm1, loss_t = train_state['val_loss'][-2:]\n",
    "         \n",
    "        # If loss worsened\n",
    "        if loss_t >= loss_tm1:\n",
    "            # Update step\n",
    "            train_state['early_stopping_step'] += 1\n",
    "        # Loss decreased\n",
    "        else:\n",
    "            # Save the best model\n",
    "            if loss_t < train_state['early_stopping_best_val']:\n",
    "                torch.save(model.state_dict(), train_state['model_filename'])\n",
    "                train_state['early_stopping_best_val'] = loss_t\n",
    "\n",
    "            # Reset early stopping step\n",
    "            train_state['early_stopping_step'] = 0\n",
    "\n",
    "        # Stop early ?\n",
    "        train_state['stop_early'] = \\\n",
    "            train_state['early_stopping_step'] >= args.early_stopping_criteria\n",
    "\n",
    "    return train_state\n",
    "\n",
    "def normalize_sizes(y_pred, y_true):\n",
    "    \"\"\"Normalize tensor sizes\n",
    "    \n",
    "    Args:\n",
    "        y_pred (torch.Tensor): the output of the model\n",
    "            If a 3-dimensional tensor, reshapes to a matrix\n",
    "        y_true (torch.Tensor): the target predictions\n",
    "            If a matrix, reshapes to be a vector\n",
    "    \"\"\"\n",
    "    if len(y_pred.size()) == 3:\n",
    "        y_pred = y_pred.contiguous().view(-1, y_pred.size(2))\n",
    "    if len(y_true.size()) == 2:\n",
    "        y_true = y_true.contiguous().view(-1)\n",
    "    return y_pred, y_true\n",
    "\n",
    "def compute_accuracy(y_pred, y_true, mask_index):\n",
    "    y_pred, y_true = normalize_sizes(y_pred, y_true)\n",
    "\n",
    "    _, y_pred_indices = y_pred.max(dim=1)\n",
    "    \n",
    "    correct_indices = torch.eq(y_pred_indices, y_true).float()\n",
    "    valid_indices = torch.ne(y_true, mask_index).float()\n",
    "    \n",
    "    n_correct = (correct_indices * valid_indices).sum().item()\n",
    "    n_valid = valid_indices.sum().item()\n",
    "\n",
    "    return n_correct / n_valid * 100\n",
    "\n",
    "def sequence_loss(y_pred, y_true, mask_index):\n",
    "    y_pred, y_true = normalize_sizes(y_pred, y_true)\n",
    "    return F.cross_entropy(y_pred, y_true, ignore_index=mask_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### General utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_seed_everywhere(seed, cuda):\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if cuda:\n",
    "        torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "def handle_dirs(dirpath):\n",
    "    if not os.path.exists(dirpath):\n",
    "        os.makedirs(dirpath)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings and some prep work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expanded filepaths: \n",
      "\tmodel_storage/ch7/model2_conditioned_surname_generation/vectorizer.json\n",
      "\tmodel_storage/ch7/model2_conditioned_surname_generation/model.pth\n",
      "Using CUDA: True\n"
     ]
    }
   ],
   "source": [
    "args = Namespace(\n",
    "    # Data and Path information\n",
    "    surname_csv=\"data/surnames/surnames_with_splits.csv\",\n",
    "    vectorizer_file=\"vectorizer.json\",\n",
    "    model_state_file=\"model.pth\",\n",
    "    save_dir=\"model_storage/ch7/model2_conditioned_surname_generation\",\n",
    "    # Model hyper parameters\n",
    "    char_embedding_size=32,\n",
    "    rnn_hidden_size=32,\n",
    "    # Training hyper parameters\n",
    "    seed=1337,\n",
    "    learning_rate=0.001,\n",
    "    batch_size=128,\n",
    "    num_epochs=100,\n",
    "    early_stopping_criteria=5,\n",
    "    # Runtime options\n",
    "    catch_keyboard_interrupt=True,\n",
    "    cuda=True,\n",
    "    expand_filepaths_to_save_dir=True,\n",
    "    reload_from_files=False,\n",
    ")\n",
    "\n",
    "if args.expand_filepaths_to_save_dir:\n",
    "    args.vectorizer_file = os.path.join(args.save_dir,\n",
    "                                        args.vectorizer_file)\n",
    "\n",
    "    args.model_state_file = os.path.join(args.save_dir,\n",
    "                                         args.model_state_file)\n",
    "    \n",
    "    print(\"Expanded filepaths: \")\n",
    "    print(\"\\t{}\".format(args.vectorizer_file))\n",
    "    print(\"\\t{}\".format(args.model_state_file))\n",
    "    \n",
    "# Check CUDA\n",
    "if not torch.cuda.is_available():\n",
    "    args.cuda = False\n",
    "\n",
    "args.device = torch.device(\"cuda\" if args.cuda else \"cpu\")\n",
    "    \n",
    "print(\"Using CUDA: {}\".format(args.cuda))\n",
    "\n",
    "# Set seed for reproducibility\n",
    "set_seed_everywhere(args.seed, args.cuda)\n",
    "\n",
    "# handle dirs\n",
    "handle_dirs(args.save_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.reload_from_files:\n",
    "    # training from a checkpoint\n",
    "    dataset = SurnameDataset.load_dataset_and_load_vectorizer(args.surname_csv,\n",
    "                                                              args.vectorizer_file)\n",
    "else:\n",
    "    # create dataset and vectorizer\n",
    "    dataset = SurnameDataset.load_dataset_and_make_vectorizer(args.surname_csv)\n",
    "    dataset.save_vectorizer(args.vectorizer_file)\n",
    "\n",
    "vectorizer = dataset.get_vectorizer()\n",
    "\n",
    "model = SurnameGenerationModel(char_embedding_size=args.char_embedding_size,\n",
    "                               char_vocab_size=len(vectorizer.char_vocab),\n",
    "                               num_nationalities=len(vectorizer.nationality_vocab),\n",
    "                               rnn_hidden_size=args.rnn_hidden_size,\n",
    "                               padding_idx=vectorizer.char_vocab.mask_index,\n",
    "                               dropout_p=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "code_folding": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbb3cb85e9464c0dab50eb822fd5e811",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='training routine', style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8de380433ae442b699a5295b3ccc5f51",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='split=train', max=60, style=ProgressStyle(description_width='…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "355a622d36894a6bbea59609014cd3fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='split=val', max=12, style=ProgressStyle(description_width='in…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exiting loop\n"
     ]
    }
   ],
   "source": [
    "mask_index = vectorizer.char_vocab.mask_index\n",
    "\n",
    "model = model.to(args.device)\n",
    "\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,\n",
    "                                           mode='min', factor=0.5,\n",
    "                                           patience=1)\n",
    "train_state = make_train_state(args)\n",
    "\n",
    "epoch_bar = tqdm_notebook(desc='training routine', \n",
    "                          total=args.num_epochs,\n",
    "                          position=0)\n",
    "\n",
    "dataset.set_split('train')\n",
    "train_bar = tqdm_notebook(desc='split=train',\n",
    "                          total=dataset.get_num_batches(args.batch_size), \n",
    "                          position=1, \n",
    "                          leave=True)\n",
    "dataset.set_split('val')\n",
    "val_bar = tqdm_notebook(desc='split=val',\n",
    "                        total=dataset.get_num_batches(args.batch_size), \n",
    "                        position=1, \n",
    "                        leave=True)\n",
    "\n",
    "try:\n",
    "    for epoch_index in range(args.num_epochs):\n",
    "        train_state['epoch_index'] = epoch_index\n",
    "\n",
    "        # Iterate over training dataset\n",
    "\n",
    "        # setup: batch generator, set loss and acc to 0, set train mode on\n",
    "        dataset.set_split('train')\n",
    "        batch_generator = generate_batches(dataset, \n",
    "                                           batch_size=args.batch_size, \n",
    "                                           device=args.device)\n",
    "        running_loss = 0.0\n",
    "        running_acc = 0.0\n",
    "        model.train()\n",
    "        \n",
    "        for batch_index, batch_dict in enumerate(batch_generator):\n",
    "            # the training routine is these 5 steps:\n",
    "\n",
    "            # --------------------------------------    \n",
    "            # step 1. zero the gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # step 2. compute the output\n",
    "            y_pred = model(x_in=batch_dict['x_data'], \n",
    "                           nationality_index=batch_dict['class_index'])\n",
    "\n",
    "            # step 3. compute the loss\n",
    "            loss = sequence_loss(y_pred, batch_dict['y_target'], mask_index)\n",
    "\n",
    "\n",
    "            # step 4. use loss to produce gradients\n",
    "            loss.backward()\n",
    "\n",
    "            # step 5. use optimizer to take gradient step\n",
    "            optimizer.step()\n",
    "            # -----------------------------------------\n",
    "            # compute the  running loss and running accuracy\n",
    "            running_loss += (loss.item() - running_loss) / (batch_index + 1)\n",
    "            acc_t = compute_accuracy(y_pred, batch_dict['y_target'], mask_index)\n",
    "            running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "\n",
    "            # update bar\n",
    "            train_bar.set_postfix(loss=running_loss,\n",
    "                                  acc=running_acc,\n",
    "                                  epoch=epoch_index)\n",
    "            train_bar.update()\n",
    "\n",
    "        train_state['train_loss'].append(running_loss)\n",
    "        train_state['train_acc'].append(running_acc)\n",
    "\n",
    "        # Iterate over val dataset\n",
    "\n",
    "        # setup: batch generator, set loss and acc to 0; set eval mode on\n",
    "        dataset.set_split('val')\n",
    "        batch_generator = generate_batches(dataset, \n",
    "                                           batch_size=args.batch_size, \n",
    "                                           device=args.device)\n",
    "        running_loss = 0.\n",
    "        running_acc = 0.\n",
    "        model.eval()\n",
    "\n",
    "        for batch_index, batch_dict in enumerate(batch_generator):\n",
    "            # compute the output\n",
    "            y_pred = model(x_in=batch_dict['x_data'], \n",
    "                           nationality_index=batch_dict['class_index'])\n",
    "\n",
    "            # step 3. compute the loss\n",
    "            loss = sequence_loss(y_pred, batch_dict['y_target'], mask_index)\n",
    "\n",
    "            # compute the  running loss and running accuracy\n",
    "            running_loss += (loss.item() - running_loss) / (batch_index + 1)\n",
    "            acc_t = compute_accuracy(y_pred, batch_dict['y_target'], mask_index)\n",
    "            running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "            \n",
    "            # Update bar\n",
    "            val_bar.set_postfix(loss=running_loss, acc=running_acc, \n",
    "                            epoch=epoch_index)\n",
    "            val_bar.update()\n",
    "\n",
    "        train_state['val_loss'].append(running_loss)\n",
    "        train_state['val_acc'].append(running_acc)\n",
    "\n",
    "        train_state = update_train_state(args=args, model=model, \n",
    "                                         train_state=train_state)\n",
    "\n",
    "        scheduler.step(train_state['val_loss'][-1])\n",
    "\n",
    "        if train_state['stop_early']:\n",
    "            break\n",
    "            \n",
    "        # move model to cpu for sampling\n",
    "        \n",
    "        nationalities = np.random.choice(np.arange(len(vectorizer.nationality_vocab)), replace=True, size=2)\n",
    "        model = model.cpu()\n",
    "        sampled_surnames = decode_samples(\n",
    "            sample_from_model(model, vectorizer, nationalities=nationalities), \n",
    "            vectorizer)\n",
    "        \n",
    "        sample1 = \"{}->{}\".format(vectorizer.nationality_vocab.lookup_index(nationalities[0]), \n",
    "                                  sampled_surnames[0])\n",
    "        sample2 = \"{}->{}\".format(vectorizer.nationality_vocab.lookup_index(nationalities[1]), \n",
    "                                  sampled_surnames[1])\n",
    "        epoch_bar.set_postfix(sample1=sample1, \n",
    "                              sample2=sample2)\n",
    "        # move model back to whichever device it should be on\n",
    "        model = model.to(args.device)\n",
    "        \n",
    "        train_bar.n = 0\n",
    "        val_bar.n = 0\n",
    "        epoch_bar.update()\n",
    "        \n",
    "except KeyboardInterrupt:\n",
    "    print(\"Exiting loop\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the loss & accuracy on the test set using the best available model\n",
    "\n",
    "model.load_state_dict(torch.load(train_state['model_filename']))\n",
    "\n",
    "model = model.to(args.device)\n",
    "\n",
    "dataset.set_split('test')\n",
    "batch_generator = generate_batches(dataset, \n",
    "                                   batch_size=args.batch_size, \n",
    "                                   device=args.device)\n",
    "running_acc = 0.\n",
    "model.eval()\n",
    "\n",
    "for batch_index, batch_dict in enumerate(batch_generator):\n",
    "    # compute the output\n",
    "    y_pred = model(x_in=batch_dict['x_data'], \n",
    "                   nationality_index=batch_dict['class_index'])\n",
    "\n",
    "    # compute the loss\n",
    "    loss = sequence_loss(y_pred, batch_dict['y_target'], mask_index)\n",
    "    \n",
    "    # compute the running loss and running accuracy\n",
    "    running_loss += (loss.item() - running_loss) / (batch_index + 1)\n",
    "    acc_t = compute_accuracy(y_pred, batch_dict['y_target'], mask_index)\n",
    "    running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "\n",
    "train_state['test_loss'] = running_loss \n",
    "train_state['test_acc'] = running_acc "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 2.480262239774068;\n",
      "Test Accuracy: 27.511795811171293\n"
     ]
    }
   ],
   "source": [
    "print(\"Test loss: {};\".format(train_state['test_loss']))\n",
    "print(\"Test Accuracy: {}\".format(train_state['test_acc']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sampled for Arabic: \n",
      "-  Gar\n",
      "-  Mamak\n",
      "-  Tamo\n",
      "Sampled for Chinese: \n",
      "-  Beag\n",
      "-  Ste\n",
      "-  Tan\n",
      "Sampled for Czech: \n",
      "-  Barkad\n",
      "-  Brimollev\n",
      "-  Maitis\n",
      "Sampled for Dutch: \n",
      "-  Smalel\n",
      "-  Dovas\n",
      "-  Tirs\n",
      "Sampled for English: \n",
      "-  Hiln\n",
      "-  Deli\n",
      "-  Kolim\n",
      "Sampled for French: \n",
      "-  Ceran\n",
      "-  Lasonis\n",
      "-  Kirteen\n",
      "Sampled for German: \n",
      "-  Pascher\n",
      "-  Sheris\n",
      "-  Bars\n",
      "Sampled for Greek: \n",
      "-  Naisastor\n",
      "-  Bonlieine\n",
      "-  Hitanani\n",
      "Sampled for Irish: \n",
      "-  Luánto\n",
      "-  Nutale\n",
      "-  Daip\n",
      "Sampled for Italian: \n",
      "-  Niliti\n",
      "-  Armane\n",
      "-  Posi\n",
      "Sampled for Japanese: \n",
      "-  Wimato\n",
      "-  Satorak\n",
      "-  Shatishi\n",
      "Sampled for Korean: \n",
      "-  Vog\n",
      "-  Ch\n",
      "-  Sta\n",
      "Sampled for Polish: \n",
      "-  Ardbanis\n",
      "-  Emitdo\n",
      "-  Altillad\n",
      "Sampled for Portuguese: \n",
      "-  Berte\n",
      "-  Balras\n",
      "-  Tenta\n",
      "Sampled for Russian: \n",
      "-  Gadekok\n",
      "-  Nanamov\n",
      "-  Zanev\n",
      "Sampled for Scottish: \n",
      "-  Pirer\n",
      "-  Serignern\n",
      "-  Sarg\n",
      "Sampled for Spanish: \n",
      "-  Carcha\n",
      "-  Aritrele\n",
      "-  Bariie\n",
      "Sampled for Vietnamese: \n",
      "-  Con\n",
      "-  Eag\n",
      "-  Di\n"
     ]
    }
   ],
   "source": [
    "model = model.cpu()\n",
    "for index in range(len(vectorizer.nationality_vocab)):\n",
    "    nationality = vectorizer.nationality_vocab.lookup_index(index)\n",
    "    print(\"Sampled for {}: \".format(nationality))\n",
    "    sampled_indices = sample_from_model(model, vectorizer,  \n",
    "                                        nationalities=[index] * 3, \n",
    "                                        temperature=0.7)\n",
    "    for sampled_surname in decode_samples(sampled_indices, vectorizer):\n",
    "        print(\"-  \" + sampled_surname)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlpbook",
   "language": "python",
   "name": "nlpbook"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "264px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "threshold": "5",
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
